#include "EndpointStore.h"
#include <algorithm>
#include <iostream>

// 外部接口声明
extern "C" {
hccl::HcclResult HcommEndPointCreate(const hccl::EndPoint *endPoint, hixl::EndpointHandle *endPointHandle);

hccl::HcclResult HcommEndPointDestroy(hixl::EndpointHandle endPointHandle);

hccl::HcclResult HcommChannelCreate(hixl::EndpointHandle *endPointHandle,
                                    uint32_t engine,
                                    void *channelDescList,
                                    uint32_t listNum,
                                    const void **memHandleList,
                                    uint32_t memHandleListNum,
                                    hixl::ChannelHandle *channelList);

hccl::HcclResult HcommChannelDestroy(const hixl::ChannelHandle *channelList, uint32_t listNum);

hccl::HcclResult HcommMemImport(hixl::EndpointHandle endPointHandle,
                                const void *memDesc,
                                uint32_t descLen,
                                void *outBuf);
}

namespace hixl {
    EndpointStore::EndpointStore() = default;

    EndpointStore::~EndpointStore() {
        // 自动清理所有端点和通道
        std::lock_guard<std::mutex> endpoints_lock(endpoints_mutex_);
        std::lock_guard<std::mutex> channels_lock(channels_mutex_);

        // 销毁所有通道
        for (auto &[endpoint, channels]: endpoint_channels_) {
            if (!channels.empty()) {
                HcommChannelDestroy(channels.data(), channels.size());
            }
        }

        // 销毁所有端点
        for (auto &[handle, endpoint_info]: endpoints_info_) {
            if (handle != nullptr) {
                HcommEndPointDestroy(handle);
            }
        }

        endpoints_info_.clear();
        endpoint_channels_.clear();
    }

    hccl::HcclResult EndpointStore::CreateEndpoint(const hccl::EndPoint &endpoint,
                                                   EndpointHandle *endpoint_handle) {
        // 调用外部接口创建端点
        hccl::HcclResult result = HcommEndPointCreate(&endpoint, endpoint_handle);
        if (result != hccl::HcclResult::SUCCESS) {
            return result;
        }

        // 保存端点信息
        {
            std::lock_guard<std::mutex> lock(endpoints_mutex_);
            EndpointInfo info;
            info.endpoint = endpoint;
            info.handle = *endpoint_handle;
            endpoints_info_[*endpoint_handle] = info;
        }

        return hccl::HcclResult::SUCCESS;
    }

    hccl::HcclResult EndpointStore::DestroyEndpoint(EndpointHandle endpoint_handle) {
        if (endpoint_handle == nullptr) {
            return hccl::HcclResult::E_PARA;
        }

        // 先销毁关联的所有通道
        {
            std::lock_guard<std::mutex> lock(channels_mutex_);
            auto it = endpoint_channels_.find(endpoint_handle);
            if (it != endpoint_channels_.end()) {
                if (!it->second.empty()) {
                    HcommChannelDestroy(it->second.data(), it->second.size());
                }
                endpoint_channels_.erase(it);
            }
        }

        // 销毁端点
        hccl::HcclResult result = HcommEndPointDestroy(endpoint_handle);
        if (result != hccl::HcclResult::SUCCESS) {
            return result;
        }

        // 从端点信息中移除
        {
            std::lock_guard<std::mutex> lock(endpoints_mutex_);
            endpoints_info_.erase(endpoint_handle);
        }

        return hccl::HcclResult::SUCCESS;
    }

    hccl::HcclResult
    EndpointStore::GetEndpointInfo(EndpointHandle endpoint_handle, hccl::EndPoint *endpoint_info) const {
        if (endpoint_handle == nullptr || endpoint_info == nullptr) {
            return hccl::HcclResult::E_PARA;
        }

        std::lock_guard<std::mutex> lock(endpoints_mutex_);
        auto it = endpoints_info_.find(endpoint_handle);
        if (it == endpoints_info_.end()) {
            return hccl::HcclResult::E_PARA;
        }

        *endpoint_info = it->second.endpoint;
        return hccl::HcclResult::SUCCESS;
    }

    const hccl::EndPoint *EndpointStore::FindEndpoint(EndpointHandle endpoint_handle) const {
        if (endpoint_handle == nullptr) {
            return nullptr;
        }

        std::lock_guard<std::mutex> lock(endpoints_mutex_);
        auto it = endpoints_info_.find(endpoint_handle);
        if (it == endpoints_info_.end()) {
            return nullptr;
        }

        return &it->second.endpoint;
    }

    hccl::HcclResult EndpointStore::CreateChannel(EndpointHandle endpoint_handle,
                                                  ChannelHandle *channel_handle) {
        if (endpoint_handle == nullptr || channel_handle == nullptr) {
            return hccl::HcclResult::E_PARA;
        }

        // 准备通道创建参数
        EndpointHandle endpoint_handles[] = {endpoint_handle};
        ChannelHandle channel_handles[1];

        hccl::HcclResult result = HcommChannelCreate(endpoint_handles,
                                                     0, // engine
                                                     nullptr, // channelDescList
                                                     1, // listNum
                                                     nullptr, // memHandleList
                                                     0, // memHandleListNum
                                                     channel_handles);
        if (result != hccl::HcclResult::SUCCESS) {
            return result;
        }

        *channel_handle = channel_handles[0]; {
            std::lock_guard<std::mutex> lock(channels_mutex_);
            endpoint_channels_[endpoint_handle].push_back(*channel_handle);
        }

        return hccl::HcclResult::SUCCESS;
    }

    hccl::HcclResult EndpointStore::DestroyChannel(EndpointHandle endpoint_handle,
                                                   ChannelHandle channel_handle) {
        if (endpoint_handle == nullptr || channel_handle == nullptr) {
            return hccl::HcclResult::E_PARA;
        }

        ChannelHandle channels[] = {channel_handle};
        hccl::HcclResult result = HcommChannelDestroy(channels, 1);
        if (result != hccl::HcclResult::SUCCESS) {
            return result;
        }

        std::lock_guard<std::mutex> lock(channels_mutex_);
        auto it = endpoint_channels_.find(endpoint_handle);
        if (it != endpoint_channels_.end()) {
            auto &channels = it->second;
            auto channel_it = std::find(channels.begin(), channels.end(), channel_handle);
            if (channel_it != channels.end()) {
                channels.erase(channel_it);
            }

            if (channels.empty()) {
                endpoint_channels_.erase(it);
            }
        }

        return hccl::HcclResult::SUCCESS;
    }

    hccl::HcclResult EndpointStore::MemImport(EndpointHandle endpoint_handle,
                                              const void *memDesc,
                                              uint32_t descLen,
                                              hccl::HcommBuf *outBuf) {
        if (endpoint_handle == nullptr || memDesc == nullptr || outBuf == nullptr) {
            return hccl::HcclResult::E_PARA;
        }

        if (descLen == 0) {
            return hccl::HcclResult::E_PARA;
        }

        hccl::HcclResult result = HcommMemImport(endpoint_handle, memDesc, descLen, outBuf);
        if (result != hccl::HcclResult::SUCCESS) {
            return result;
        }

        return hccl::HcclResult::SUCCESS;
    }

    std::vector<EndpointHandle> EndpointStore::GetAllEndpointHandles() const {
        std::lock_guard<std::mutex> lock(endpoints_mutex_);
        std::vector<EndpointHandle> handles;
        for (const auto &[handle, info]: endpoints_info_) {
            handles.push_back(handle);
        }
        return handles;
    }
} // namespace hixl
